import os
import torch

def find_desirable_params(state_dict, org_prefix, desirable_prefix):
    if not org_prefix.endswith("."):
        org_prefix = org_prefix + "."
    # if not desirable_prefix.endswith("."):
    #     desirable_prefix = desirable_prefix + "."
    state_interest = {}
    for k, v in state_dict.items():
        if k.startswith(org_prefix):
            if desirable_prefix == "":
                state_interest[k[len(org_prefix): ]] = v
            else:
                state_interest[desirable_prefix + k[len(org_prefix): ]] = v
    if len(state_interest) == 0:    
        raise NotImplementedError(f"Prefix {org_prefix} not exist. A referred prefix: {k}")
    return state_interest

def model_init(model: torch.nn.Module, ckpt, prefix: (str or list), desirable_prefix: (str or list)):
    if ckpt is None:
        return model
    print("Loading ckpt:", ckpt)
    state_dict = torch.load(ckpt, map_location="cpu")["state_dict"]
    if isinstance(prefix, str):
        prefix = [prefix]
    if isinstance(desirable_prefix, str):
        desirable_prefix = [desirable_prefix]
    assert len(prefix) == len(desirable_prefix), "prefix and desirable_prefix should have same length."
    state_all = {}
    for i in range(len(prefix)):
        print("Substitute prefix:", prefix[i], "\t\t with desirable_prefix:", desirable_prefix[i])
        state_interest = find_desirable_params(state_dict, prefix[i], desirable_prefix[i])
        state_all.update(state_interest)
    missing, unexpected = model.load_state_dict(state_all, strict=False)
    print("missing:", missing)
    print("unexpected:", unexpected)
    return model